[FlyDSL] Read (B, next_n) context lengths in the paged MQA logits kernel - #6450

Motivation

Currently, vLLM passes context_lens to the paged MQA logits kernels as a 2D (B, next_n) table, the same layout DeepGEMM uses. The FlyDSL kernel only reads a 1D [B] vector, so with MTP on, flydsl_fp8_paged_mqa_logits fails with a reshape error:

RuntimeError: shape '[4]' is invalid for input of size 12

This PR lets the FlyDSL paged_mqa_logits kernel accept the 2D context_lens that vLLM sends with MTP, while staying backward compatible with the 1D format.

Split out of #6153 as requested in review (one backend per PR); the Triton and Gluon change is in #6451. The two PRs touch different files and can merge in either order.

Example: 4 requests, next_n=3, context lengths 8192 / 1500 / 3000 / 700, vLLM sends the following 2D array:

The fix uses the row length (next_n) as the stride and reads the last entry of each row, which gives [8192, 1500, 3000, 700].

What changes

• The stride is a build-time constant. The kernel reads request b's length at b * ctx_stride + ctx_stride - 1; stride 1 builds the pre-PR kernel (same code and name), stride n builds a variant named ..._cs{n}, and _compile() caches one build per stride.
  • 1D [B]: stride is 1, so this is context_lens[b], the same load as today.
  • 2D (B, n): stride is n, so this is the last entry of row b, which is the full sequence length.
• The wrapper picks the stride from the tensor's rank, and raises ValueError unless context_lens is a contiguous int32 [B] or (B, n) on the same device as q_fp8, with exactly B rows and n equal to 1 or next_n:
  • a non-contiguous input is rejected rather than copied, so there is no hidden allocation per call. Before, reshape silently copied some strided inputs and kept others as strided views that the kernel read as dense.
  • a (B, 0) table would have made the kernel read before the start of the tensor, and a (B, n) with n other than next_n read the wrong column as the length.
  • a wrong row count raised a RuntimeError from reshape before; it is now a ValueError with the expected shape.
  • the kernel reads context_lens as int32 words, so an int64 table gave wrong lengths before.

Backward compatibility

• A contiguous [B] is unchanged, and (B, 1) behaves exactly like [B].
• A contiguous (B, next_n) table is the only new input. vLLM already keeps its decode seq_lens contiguous for this.
• A strided, non-int32 or off-device context_lens now raises ValueError.

Tests

In op_tests/test_flydsl_fp8_paged_mqa_logits.py:

• The context_lens layout is a new axis of the existing sweep, -c/--ctx_lens (default 1d 2d). For 2d, FlyDSL gets the (B, next_n) table vLLM builds, is graded against the same fp32 reference as 1d, and is timed next to it in the summary table.
• test_fp8_paged_mqa_logits_mixed_context_lens: B=5, next_n=3, distinct lengths and one empty (CUDA-graph padding) request, graded against the fp32 reference for 1d and 2d, with and without SplitKV. It fails if the kernel reads row 0's length for every request.
• test_fp8_paged_mqa_logits_bad_context_lens: too few or too many rows, no columns, a column count other than next_n, an int64 table, a CPU table, a strided [B] and a non-contiguous (B, n) must raise ValueError.

Results

MI355X, Triton 3.8, FlyDSL 0.3.4.1, upstream main at 9fd025efd plus this change:

• python op_tests/test_flydsl_fp8_paged_mqa_logits.py: all tables pass. In the default sweep (batch 1-8, next_n 1-8, 1k context) 2d has the same error against the reference as 1d in every case, and its kernel time is within +0.1% to +1.9% of 1d.
• The new test fails on main without this change, with the reshape error above.
• black and ruff (0.16.0) are clean.
• Kernel time at the GLM-5 indexer shape with 8k context: 1D matches the kernel on main (e.g. B=16, next_n=3: 30.6 / 30.2 us vs 30.7 / 30.2 us over two runs), and 2D stays within about 1% of 1D.